#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
逐帧跳绳分析器 - 提取MediaPipe数据到CSV文件
基于手部动作、脚部动作和人物重心移动的精确跳绳检测
"""

import cv2
import mediapipe as mp
import numpy as np
import pandas as pd
import os
from typing import List, Dict, Tuple, Optional
import argparse
from pathlib import Path

class FrameByFrameJumpRopeAnalyzer:
    def __init__(self):
        # 初始化MediaPipe
        self.mp_pose = mp.solutions.pose
        self.pose = self.mp_pose.Pose(
            static_image_mode=False,
            model_complexity=2,
            enable_segmentation=False,
            min_detection_confidence=0.7,
            min_tracking_confidence=0.5
        )
        self.mp_drawing = mp.solutions.drawing_utils
        
        # 关键点索引定义
        self.POSE_LANDMARKS = {
            # 头部
            'nose': 0,
            'left_eye_inner': 1, 'left_eye': 2, 'left_eye_outer': 3,
            'right_eye_inner': 4, 'right_eye': 5, 'right_eye_outer': 6,
            'left_ear': 7, 'right_ear': 8,
            'mouth_left': 9, 'mouth_right': 10,
            
            # 上身
            'left_shoulder': 11, 'right_shoulder': 12,
            'left_elbow': 13, 'right_elbow': 14,
            'left_wrist': 15, 'right_wrist': 16,
            'left_pinky': 17, 'right_pinky': 18,
            'left_index': 19, 'right_index': 20,
            'left_thumb': 21, 'right_thumb': 22,
            
            # 下身
            'left_hip': 23, 'right_hip': 24,
            'left_knee': 25, 'right_knee': 26,
            'left_ankle': 27, 'right_ankle': 28,
            'left_heel': 29, 'right_heel': 30,
            'left_foot_index': 31, 'right_foot_index': 32
        }
        
    def extract_landmarks_to_csv(self, video_path: str, output_csv: str) -> bool:
        """
        逐帧提取MediaPipe关键点数据到CSV文件
        """
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            print(f"错误：无法打开视频文件 {video_path}")
            return False
            
        # 获取视频信息
        fps = cap.get(cv2.CAP_PROP_FPS)
        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        duration = total_frames / fps
        
        print(f"视频信息：")
        print(f"  文件: {video_path}")
        print(f"  帧率: {fps} FPS")
        print(f"  总帧数: {total_frames}")
        print(f"  时长: {duration:.2f} 秒")
        
        # 准备CSV数据结构
        csv_data = []
        
        frame_count = 0
        while True:
            ret, frame = cap.read()
            if not ret:
                break
                
            # 转换为RGB
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            
            # MediaPipe姿态检测
            results = self.pose.process(rgb_frame)
            
            # 当前帧的数据行
            row_data = {
                'frame_number': frame_count,
                'timestamp': frame_count / fps,
                'has_pose': results.pose_landmarks is not None
            }
            
            if results.pose_landmarks:
                # 提取所有关键点坐标
                landmarks = results.pose_landmarks.landmark
                
                for name, idx in self.POSE_LANDMARKS.items():
                    if idx < len(landmarks):
                        landmark = landmarks[idx]
                        row_data[f'{name}_x'] = landmark.x
                        row_data[f'{name}_y'] = landmark.y
                        row_data[f'{name}_z'] = landmark.z
                        row_data[f'{name}_visibility'] = landmark.visibility
                    else:
                        row_data[f'{name}_x'] = None
                        row_data[f'{name}_y'] = None
                        row_data[f'{name}_z'] = None
                        row_data[f'{name}_visibility'] = None
                        
                # 计算关键特征
                self._calculate_frame_features(row_data, landmarks)
            else:
                # 没有检测到姿态时填充空值
                for name in self.POSE_LANDMARKS.keys():
                    row_data[f'{name}_x'] = None
                    row_data[f'{name}_y'] = None
                    row_data[f'{name}_z'] = None
                    row_data[f'{name}_visibility'] = None
                    
                # 特征也设为空
                self._set_empty_features(row_data)
            
            csv_data.append(row_data)
            frame_count += 1
            
            # 进度显示
            if frame_count % 100 == 0:
                progress = (frame_count / total_frames) * 100
                print(f"处理进度: {progress:.1f}% ({frame_count}/{total_frames})")
        
        cap.release()
        
        # 保存到CSV
        df = pd.DataFrame(csv_data)
        df.to_csv(output_csv, index=False, encoding='utf-8-sig')
        print(f"\n数据已保存到: {output_csv}")
        print(f"总共处理了 {frame_count} 帧")
        
        return True
    
    def _calculate_frame_features(self, row_data: Dict, landmarks) -> None:
        """
        计算当前帧的关键特征 (增强版)
        """
        try:
            # 提前获取所有需要的关键点，避免重复索引
            left_shoulder = landmarks[self.POSE_LANDMARKS['left_shoulder']]
            right_shoulder = landmarks[self.POSE_LANDMARKS['right_shoulder']]
            left_hip = landmarks[self.POSE_LANDMARKS['left_hip']]
            right_hip = landmarks[self.POSE_LANDMARKS['right_hip']]
            left_knee = landmarks[self.POSE_LANDMARKS['left_knee']]
            right_knee = landmarks[self.POSE_LANDMARKS['right_knee']]
            left_ankle = landmarks[self.POSE_LANDMARKS['left_ankle']]
            right_ankle = landmarks[self.POSE_LANDMARKS['right_ankle']]
            left_wrist = landmarks[self.POSE_LANDMARKS['left_wrist']]
            right_wrist = landmarks[self.POSE_LANDMARKS['right_wrist']]
            left_elbow = landmarks[self.POSE_LANDMARKS['left_elbow']]
            right_elbow = landmarks[self.POSE_LANDMARKS['right_elbow']]

            # 1. 重心位置 (髋部中点)
            center_of_mass_y = (left_hip.y + right_hip.y) / 2
            row_data['center_of_mass_y'] = center_of_mass_y
            
            # 2. 脚踝高度和置信度 (修改: 分开记录)
            row_data['left_ankle_y'] = left_ankle.y
            row_data['right_ankle_y'] = right_ankle.y
            row_data['left_ankle_visibility'] = left_ankle.visibility
            row_data['right_ankle_visibility'] = right_ankle.visibility
            # 保留平均值以兼容旧算法，后续可移除
            avg_ankle_y = (left_ankle.y + right_ankle.y) / 2
            row_data['avg_ankle_y'] = avg_ankle_y
            
            # 3. 手腕位置和置信度 (修改: 补充置信度)
            row_data['left_wrist_y'] = left_wrist.y
            row_data['right_wrist_y'] = right_wrist.y
            row_data['left_wrist_visibility'] = left_wrist.visibility
            row_data['right_wrist_visibility'] = right_wrist.visibility
            # 保留平均值
            row_data['avg_wrist_y'] = (left_wrist.y + right_wrist.y) / 2
            
            # 4. 肩膀位置
            avg_shoulder_y = (left_shoulder.y + right_shoulder.y) / 2
            row_data['avg_shoulder_y'] = avg_shoulder_y
            
            # 5. 身体高度 (肩膀到脚踝的距离)
            body_height = avg_ankle_y - avg_shoulder_y
            row_data['body_height'] = body_height
            
            # 6. 手臂角度 (肘部角度)
            left_arm_angle = self._calculate_angle([left_shoulder.x, left_shoulder.y], [left_elbow.x, left_elbow.y], [left_wrist.x, left_wrist.y])
            row_data['left_arm_angle'] = left_arm_angle
            right_arm_angle = self._calculate_angle([right_shoulder.x, right_shoulder.y], [right_elbow.x, right_elbow.y], [right_wrist.x, right_wrist.y])
            row_data['right_arm_angle'] = right_arm_angle
            
            # 7. 腿部角度 (膝盖角度)
            left_leg_angle = self._calculate_angle([left_hip.x, left_hip.y], [left_knee.x, left_knee.y], [left_ankle.x, left_ankle.y])
            row_data['left_leg_angle'] = left_leg_angle
            right_leg_angle = self._calculate_angle([right_hip.x, right_hip.y], [right_knee.x, right_knee.y], [right_ankle.x, right_ankle.y])
            row_data['right_leg_angle'] = right_leg_angle
            
            # 8. 手腕相对于肩膀的位置
            row_data['left_wrist_relative_to_shoulder'] = left_wrist.y - left_shoulder.y
            row_data['right_wrist_relative_to_shoulder'] = right_wrist.y - right_shoulder.y
            
            # --- 以下为新增特征 ---

            # 9. 身体姿态的包围盒 (Bounding Box) 特征
            visible_landmarks_x = [lm.x for lm in landmarks if lm.visibility > 0.5]
            visible_landmarks_y = [lm.y for lm in landmarks if lm.visibility > 0.5]

            if visible_landmarks_x and visible_landmarks_y:
                min_x, max_x = min(visible_landmarks_x), max(visible_landmarks_x)
                min_y, max_y = min(visible_landmarks_y), max(visible_landmarks_y)
                row_data['pose_bbox_width'] = max_x - min_x
                row_data['pose_bbox_height'] = max_y - min_y
                row_data['pose_aspect_ratio'] = (max_y - min_y) / (max_x - min_x) if (max_x - min_x) > 0 else 0
            else:
                row_data['pose_bbox_width'] = None
                row_data['pose_bbox_height'] = None
                row_data['pose_aspect_ratio'] = None

            # 10. 姿态对称性和稳定性特征
            row_data['shoulder_tilt'] = abs(left_shoulder.y - right_shoulder.y)
            row_data['hip_tilt'] = abs(left_hip.y - right_hip.y)
            row_data['left_wrist_hip_dy'] = left_wrist.y - left_hip.y
            row_data['right_wrist_hip_dy'] = right_wrist.y - right_hip.y
            row_data['ankle_distance'] = np.sqrt((left_ankle.x - right_ankle.x)**2 + (left_ankle.y - right_ankle.y)**2)
            
            # 11. 身体前倾角度
            shoulder_mid_x = (left_shoulder.x + right_shoulder.x) / 2
            shoulder_mid_y = (left_shoulder.y + right_shoulder.y) / 2
            hip_mid_x = (left_hip.x + right_hip.x) / 2
            hip_mid_y = (left_hip.y + right_hip.y) / 2

            torso_vector = np.array([hip_mid_x - shoulder_mid_x, hip_mid_y - shoulder_mid_y])
            vertical_vector = np.array([0, -1]) # 垂直向上向量

            norm_torso = np.linalg.norm(torso_vector)
            if norm_torso > 0:
                cos_angle = np.dot(torso_vector, vertical_vector) / norm_torso
                angle_rad = np.arccos(np.clip(cos_angle, -1.0, 1.0))
                row_data['body_lean_angle'] = np.degrees(angle_rad)
            else:
                row_data['body_lean_angle'] = None
            
        except Exception as e:
            print(f"帧 {row_data.get('frame_number', 'N/A')} 计算特征时出错: {e}")
            self._set_empty_features(row_data)
    
    def _set_empty_features(self, row_data: Dict) -> None:
        """
        设置空特征值 (增强版)
        """
        # 包含所有旧特征和新特征
        features = [
            # 旧有特征
            'center_of_mass_y', 'avg_ankle_y', 'left_wrist_y', 'right_wrist_y',
            'avg_wrist_y', 'avg_shoulder_y', 'body_height', 'left_arm_angle',
            'right_arm_angle', 'left_leg_angle', 'right_leg_angle',
            'left_wrist_relative_to_shoulder', 'right_wrist_relative_to_shoulder',
            
            # 新增和修改的特征
            'left_ankle_y', 'right_ankle_y', 'left_ankle_visibility', 'right_ankle_visibility',
            'left_wrist_visibility', 'right_wrist_visibility',
            'pose_bbox_width', 'pose_bbox_height', 'pose_aspect_ratio',
            'shoulder_tilt', 'hip_tilt', 'left_wrist_hip_dy', 'right_wrist_hip_dy',
            'ankle_distance', 'body_lean_angle'
        ]
        for feature in features:
            # 确保即使在计算中途出错，该字段也存在于row_data中
            if feature not in row_data:
                row_data[feature] = None
    
    def _calculate_angle(self, a: List[float], b: List[float], c: List[float]) -> float:
        """
        计算三点之间的角度
        """
        try:
            a = np.array(a)
            b = np.array(b)
            c = np.array(c)
            
            radians = np.arctan2(c[1] - b[1], c[0] - b[0]) - np.arctan2(a[1] - b[1], a[0] - b[0])
            angle = np.abs(radians * 180.0 / np.pi)
            
            if angle > 180.0:
                angle = 360 - angle
                
            return angle
        except:
            return None

def main():
    parser = argparse.ArgumentParser(description='逐帧跳绳分析器')
    parser.add_argument('--video', required=True, help='输入视频文件路径')
    parser.add_argument('--output', help='输出CSV文件路径（可选）')
    
    args = parser.parse_args()
    
    video_path = args.video
    if not os.path.exists(video_path):
        print(f"错误：视频文件不存在 {video_path}")
        return
    
    # 生成输出文件名
    if args.output:
        output_csv = args.output
    else:
        video_name = Path(video_path).stem
        output_csv = f"{video_name}_frame_data.csv"
    
    # 创建分析器并处理
    analyzer = FrameByFrameJumpRopeAnalyzer()
    success = analyzer.extract_landmarks_to_csv(video_path, output_csv)
    
    if success:
        print(f"\n✅ 成功完成逐帧分析！")
        print(f"CSV文件已保存到: {output_csv}")
        print(f"\n接下来可以使用Excel或其他工具打开CSV文件进行分析")
        print(f"主要分析指标：")
        print(f"  - center_of_mass_y: 人体重心Y坐标")
        print(f"  - avg_ankle_y: 脚踝平均Y坐标")
        print(f"  - avg_wrist_y: 手腕平均Y坐标")
        print(f"  - left/right_arm_angle: 左右臂角度")
        print(f"  - left/right_leg_angle: 左右腿角度")
        print(f"  - body_lean_angle: 身体前倾角度")
    else:
        print("❌ 分析失败")

if __name__ == "__main__":
    main()